"""
机器学习模块 - 负责训练和优化AI模型

该模块负责训练强化学习模型，管理训练过程，
保存和加载模型权重，以及模型性能评估。
"""

import numpy as np
import tensorflow as tf
from tensorflow import keras
import os
import json
from typing import Dict, List, Tuple, Any
from datetime import datetime


class ModelTrainer:
    """
    模型训练类 - 训练和优化AI模型
    
    该类负责：
    1. 管理训练过程
    2. 保存和加载模型
    3. 评估模型性能
    """
    
    def __init__(self, model_save_dir: str = "data/models"):
        """
        初始化模型训练器
        
        Args:
            model_save_dir: 模型保存目录
        """
        self.model_save_dir = model_save_dir
        
        # 创建模型保存目录
        if not os.path.exists(model_save_dir):
            os.makedirs(model_save_dir)
        
        # 训练历史记录
        self.training_history = []
        
        print("模型训练器初始化完成")
    
    def save_training_data(self, data: Dict[str, Any], filename: str = None) -> str:
        """
        保存训练数据
        
        Args:
            data: 训练数据
            filename: 文件名（可选）
            
        Returns:
            str: 保存的文件路径
        """
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"training_data_{timestamp}.json"
        
        filepath = os.path.join(self.model_save_dir, filename)
        
        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            print(f"训练数据已保存至: {filepath}")
            return filepath
        except Exception as e:
            print(f"训练数据保存失败: {e}")
            return ""
    
    def load_training_data(self, filepath: str) -> Dict[str, Any]:
        """
        加载训练数据
        
        Args:
            filepath: 数据文件路径
            
        Returns:
            Dict[str, Any]: 加载的训练数据
        """
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)
            print(f"训练数据已从 {filepath} 加载")
            return data
        except Exception as e:
            print(f"训练数据加载失败: {e}")
            return {}
    
    def evaluate_model(self, model: keras.Model, test_data: Tuple[np.ndarray, np.ndarray]) -> Dict[str, float]:
        """
        评估模型性能
        
        Args:
            model: 要评估的模型
            test_data: 测试数据 (features, labels)
            
        Returns:
            Dict[str, float]: 评估结果
        """
        try:
            # 评估模型
            loss = model.evaluate(test_data[0], test_data[1], verbose=0)
            
            # 计算准确率（示例）
            predictions = model.predict(test_data[0], verbose=0)
            predicted_classes = np.argmax(predictions, axis=1)
            true_classes = np.argmax(test_data[1], axis=1)
            accuracy = np.mean(predicted_classes == true_classes)
            
            evaluation = {
                'loss': loss,
                'accuracy': accuracy
            }
            
            print(f"模型评估结果 - Loss: {loss:.4f}, Accuracy: {accuracy:.4f}")
            return evaluation
        except Exception as e:
            print(f"模型评估失败: {e}")
            return {}
    
    def train_with_replay(self, ai_agent, episodes: int = 1000) -> None:
        """
        使用经验回放训练AI代理
        
        Args:
            ai_agent: AI决策代理
            episodes: 训练回合数
        """
        try:
            batch_size = 32
            
            for episode in range(episodes):
                # 模拟游戏过程（这里应该是与游戏环境的交互）
                # 为简化示例，我们使用随机数据
                
                # 每100回合更新目标网络
                if episode % 100 == 0:
                    ai_agent.update_target_network()
                    print(f"回合 {episode}: 目标网络已更新")
                
                # 每10回合进行经验回放训练
                if episode % 10 == 0:
                    ai_agent.replay(batch_size)
                
                # 每100回合保存模型
                if episode % 100 == 0:
                    model_path = os.path.join(
                        self.model_save_dir, 
                        f"model_episode_{episode}.h5"
                    )
                    ai_agent.save_model(model_path)
                
                # 记录训练历史
                self.training_history.append({
                    'episode': episode,
                    'epsilon': ai_agent.epsilon,
                    'memory_size': len(ai_agent.memory)
                })
                
                if episode % 100 == 0:
                    print(f"训练进度: {episode}/{episodes} 回合")
            
            # 保存最终模型
            final_model_path = os.path.join(self.model_save_dir, "final_model.h5")
            ai_agent.save_model(final_model_path)
            
            # 保存训练历史
            history_path = os.path.join(self.model_save_dir, "training_history.json")
            self.save_training_data({
                'history': self.training_history,
                'episodes': episodes
            }, "training_history.json")
            
            print("训练完成")
            
        except Exception as e:
            print(f"训练过程出错: {e}")
    
    def compare_models(self, model_paths: List[str], test_data: Tuple[np.ndarray, np.ndarray]) -> Dict[str, Dict[str, float]]:
        """
        比较多个模型的性能
        
        Args:
            model_paths: 模型路径列表
            test_data: 测试数据
            
        Returns:
            Dict[str, Dict[str, float]]: 各模型评估结果
        """
        results = {}
        
        try:
            for path in model_paths:
                if os.path.exists(path):
                    # 创建新模型并加载权重
                    # 注意：这里需要根据实际模型结构创建模型
                    model = keras.Sequential([
                        keras.layers.Dense(64, activation='relu'),
                        keras.layers.Dense(64, activation='relu'),
                        keras.layers.Dense(32, activation='relu'),
                        keras.layers.Dense(20, activation='linear')
                    ])
                    model.compile(optimizer='adam', loss='mse')
                    
                    model.load_weights(path)
                    
                    # 评估模型
                    evaluation = self.evaluate_model(model, test_data)
                    results[path] = evaluation
                else:
                    print(f"模型文件不存在: {path}")
            
            return results
        except Exception as e:
            print(f"模型比较失败: {e}")
            return {}


# 使用示例
if __name__ == "__main__":
    # 创建模型训练器实例
    trainer = ModelTrainer()
    
    print("模型训练器模块测试完成")